package com.ruoyi.handler.limit;

import com.ruoyi.handler.utils.ThreadPoolUtils;

import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;

/**
 * 令牌桶限流
 */
public class BucketLimiter {
    private static final Integer CAPACITY = 10;

    /**
     * 当前桶内令牌数
     */
    private static final AtomicLong CUR_CAP = new AtomicLong(10);

    private static final Integer SPEED_PER_SECOND = 2;

    /**
     * 单机限流
     *
     * @return boolean 是否被限流
     */
    public static boolean singleLimit() {

        // 桶里还有令牌，正常通过请求
        if (CUR_CAP.get() > 0) {
            CUR_CAP.decrementAndGet();
            return false;
        }

        // 没了，限流
        return true;
    }

    /**
     * 每秒往桶里增加两个令牌
     */
    public static void addBucketTask() {
        ThreadPoolUtils.scheduledTask(() -> {
            if (CUR_CAP.get() < CAPACITY) {
                CUR_CAP.addAndGet(2);
            }
        }, 1, TimeUnit.SECONDS);
    }

}
